import numpy as np
import chromadb
from sentence_transformers import SentenceTransformer
from collections import defaultdict
import json
import os
import shutil
from typing import Dict, List
import time
from sklearn.metrics.pairwise import cosine_similarity


class DataPreparationLayer:
    def __init__(
        self, 
        model_path: str = "/home/ai/huggingface_cache/BAAI/bge-large-zh-v1.5/bge-large-zh-v1.5",
        chroma_persist_dir: str = "/home/ai/Recommendation_system/chroma_template_db",
        behavior_data_path: str = "./user_behavior_data.json"
    ):
        """初始化数据准备层（适配ChromaDB最新版）"""
        
        # 确保目录存在且可写
        self._ensure_directory(chroma_persist_dir)
        
        # 初始化模型
        try:
            self.model = SentenceTransformer(model_path)
            print(f"✓ 模型加载成功：{model_path}")
        except Exception as e:
            print(f"✗ 模型加载失败：{e}")
            raise
        
        # 初始化ChromaDB客户端
        try:
            self.chroma_client = chromadb.PersistentClient(
                path=chroma_persist_dir
            )
            print(f"✓ ChromaDB客户端初始化成功")
            print(f"  存储位置：{os.path.abspath(chroma_persist_dir)}")
        except Exception as e:
            print(f"✗ ChromaDB初始化失败：{e}")
            raise
        
        # 获取或创建集合
        self.template_collection = self.chroma_client.get_or_create_collection(
            name="template_embeddings",
            metadata={"description": "存储模板描述的语义向量"}
        )


        

        

        # 初始化内存数据结构
        self.user_behavior = defaultdict(dict)  # 初始化一个嵌套字典
        self.behavior_data_path = behavior_data_path
        self.templates = {}
        
        # 加载历史数据
        self._load_behavior_data()


                
        
        # 验证数据库连接
        self._verify_db_connection()

        # 缓存相关属性
        self.user_similarities = None  # 存储用户相似度缓存
        self.template_ids = []  # 缓存所有模板ID
        self.last_similarity_update = 0  # 最后更新时间戳
        self._sync_template_ids()  # 初始化模板ID列表

    def _ensure_directory(self, path: str) -> None:
        """确保目录存在且可写"""
        try:
            os.makedirs(path, exist_ok=True)
            # 测试写入权限
            test_file = os.path.join(path, ".permission_test")
            with open(test_file, 'w') as f:
                f.write("test")
            os.remove(test_file)
            print(f"✓ 目录验证通过：{path}")
        except Exception as e:
            print(f"✗ 目录权限错误：{path} - {e}")
            raise

    def _verify_db_connection(self) -> bool:
        """验证数据库连接状态"""
        try:
            count = self.template_collection.count()
            print(f"✓ 数据库连接验证成功，当前模板数：{count}")
            return True
        except Exception as e:
            print(f"✗ 数据库连接验证失败：{e}")
            return False

    def _load_behavior_data(self) -> None:
        """加载用户行为数据"""
        if os.path.exists(self.behavior_data_path):
            try:
                with open(self.behavior_data_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    for user_id, behaviors in data.items():
                        self.user_behavior[int(user_id)] = {
                            int(tid): score for tid, score in behaviors.items()
                        }
                print(f"✓ 已加载历史用户行为数据，共{len(self.user_behavior)}个用户")
            except Exception as e:
                print(f"✗ 加载用户行为数据失败：{e}，将使用空数据")

    def save_behavior_data(self) -> bool:
        """保存用户行为数据"""
        try:
            with open(self.behavior_data_path, 'w', encoding='utf-8') as f:
                data = {
                    str(user_id): {str(tid): score for tid, score in behaviors.items()}
                    for user_id, behaviors in self.user_behavior.items()
                }
                json.dump(data, f, ensure_ascii=False, indent=2)
            print(f"✓ 用户行为数据已保存至{self.behavior_data_path}")
            return True
        except Exception as e:
            print(f"✗ 保存用户行为数据失败：{e}")
            return False

    def add_template(self, template_id: int, name: str, description: str, hangye: str) -> bool:
        """添加单个模板"""
        if template_id in self.templates:
            print(f"⚠ 模板{template_id}已存在，跳过添加")
            return False
            
        # 生成向量
        try:
            embedding = self.model.encode(
                description,
                normalize_embeddings=True,
                show_progress_bar=False
            ).tolist()
            
            if len(embedding) != 1024:
                print(f"✗ 模板{template_id}向量维度错误：预期1024维，实际{len(embedding)}维")
                return False
            print(f"✓ 模板{template_id}向量生成成功（1024维）")
        except Exception as e:
            print(f"✗ 模板{template_id}向量生成失败：{e}")
            return False
            
        # 存储模板元数据
        self.templates[template_id] = {
            "name": name,
            "description": description,
            "hangye": hangye
        }
        
        # 写入ChromaDB
        try:
            self.template_collection.add(
                ids=[str(template_id)],
                embeddings=[embedding],
                metadatas=[{
                    "name": name,
                    "template_id": template_id,
                    "description": description,
                    "hangye": hangye
                }]
            )
            print(f"✓ 模板{template_id}（{name}）成功写入数据库")
            return True
        except Exception as e:
            print(f"✗ 模板{template_id}写入失败：{e}")
            return False

    def batch_add_templates(self, templates: List[Dict]) -> int:
        """批量添加模板"""
        success_count = 0
        for tpl in templates:
            if self.add_template(tpl["id"], tpl["name"], tpl["description"], tpl["hangye"]):
                success_count += 1
        
        print(f"✓ 批量添加完成，成功添加{success_count}/{len(templates)}个模板")
        return success_count

    def record_user_behavior(self, user_id: int, template_id: int, behavior_type: str) -> bool:
        """记录用户行为"""
        if template_id not in self.templates:
            print(f"✗ 模板{template_id}不存在，无法记录行为")
            return False
            
        behavior_weights = {'view': 1, 'collect': 2, 'download': 3}
        score = behavior_weights.get(behavior_type, 1)
        current_score = self.user_behavior[user_id].get(template_id, 0)
        
        if score > current_score:
            self.user_behavior[user_id][template_id] = score
            print(f"✓ 用户{user_id}对模板{template_id}的{behavior_type}行为已记录（分数：{score}）")
            return True
        return False

    def export_template_vectors(self, output_path: str = "./template_vectors.json") -> bool:
        """导出模板向量"""
        if not self.templates:
            print("✗ 没有模板数据可导出")
            return False
            
        try:
            data = {
                "template_count": len(self.templates),
                "templates": []
            }
            
            # 直接从内存数据生成，确保一致性
            for template_id, info in self.templates.items():
                embedding = self.model.encode(
                    info["description"],
                    normalize_embeddings=True
                ).tolist()
                
                data["templates"].append({
                    "template_id": template_id,
                    "name": info["name"],
                    "hangye": info["hangye"],
                    "vector_dim": len(embedding),
                    "vector_sample": embedding[:5]  # 只取前5维作为示例
                })
                
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(data, f, ensure_ascii=False, indent=2)
            print(f"✓ 模板向量信息已导出至{output_path}")
            return True
        except Exception as e:
            print(f"✗ 导出模板向量失败：{e}")
            return False



#-----------------------------------------------------------------------------------------------------------


    def get_database_status(self) -> Dict:
        """获取数据库状态信息"""
        try:
            db_path = self.chroma_client._settings.persist_directory
            return {
                "db_path": os.path.abspath(db_path),
                "template_count": self.template_collection.count(),
                "disk_usage": f"{self._get_dir_size(db_path)/1024:.2f} KB",
                "files": os.listdir(db_path) if os.path.exists(db_path) else []
            }
        except Exception as e:
            return {"error": str(e)}

    def _get_dir_size(self, path: str) -> int:
        """计算目录大小"""
        total = 0
        for dirpath, dirnames, filenames in os.walk(path):
            for f in filenames:
                fp = os.path.join(dirpath, f)
                total += os.path.getsize(fp)
        return total

    def get_template_count(self) -> int:
        return len(self.templates)

    def get_user_count(self) -> int:
        return len(self.user_behavior)

    def get_template_info(self, template_id: int) -> Dict:
        return self.templates.get(template_id, None)
    
    # 读取 JSON 中的模板数据
    @staticmethod
    def load_templates_from_json(file_path="./sample_templates.json"):
        """从 JSON 文件加载模板数据"""
        if not os.path.exists(file_path):
            print(f"模板数据文件 {file_path} 不存在")
            return []
        
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                templates = json.load(f)
            print(f"从 {file_path} 加载模板数据成功，共 {len(templates)} 条")
            return templates
        except Exception as e:
            print(f"读取模板数据失败：{e}")
            return []

    @staticmethod
    def load_and_record_behavior(data_layer, file_path="./user_behavior_samples.json"):
        """从JSON文件加载用户行为并记录到数据层"""
        if not os.path.exists(file_path):
            print(f"用户行为数据文件 {file_path} 不存在")
            return
        
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                behaviors = json.load(f)
            
            # 遍历行为数据并记录
            for idx, behavior in enumerate(behaviors, 1):
                user_id = behavior["user_id"]
                template_id = behavior["template_id"]
                behavior_type = behavior["behavior_type"]
                
                # 调用record_user_behavior记录
                success = data_layer.record_user_behavior(
                    user_id=user_id,
                    template_id=template_id,
                    behavior_type=behavior_type
                )
                
                if success:
                    print(f"行为{idx}：用户{user_id} {behavior_type} 模板{template_id} → 记录成功")
                else:
                    print(f"行为{idx}：用户{user_id} {behavior_type} 模板{template_id} → 记录失败（可能已存在更高权重）")
        
        except Exception as e:
            print(f"加载并记录用户行为失败：{e}")


    # 模板ID同步方法
    def _sync_template_ids(self):
        """同步ChromaDB中的模板ID到本地列表"""
        try:
            all_ids = self.template_collection.get()["ids"]
            self.template_ids = [int(tid) for tid in all_ids]  # 转为整数ID
            print(f"已同步模板ID列表，共{len(self.template_ids)}个模板")
        except Exception as e:
            print(f"同步模板ID失败：{e}")
            self.template_ids = []


#--------------------------------------------------计算用户相似度并缓存到本地JSON文件---------------------------------------------------
    
    def _calculate_user_similarity(self, force_recalculate=False, expire_seconds=3600, save_path="/home/ai/Recommendation_system/user_similarities_data.json") -> dict:
        """
        计算并缓存用户相似度，并保存到本地JSON文件
        :param force_recalculate: 强制重新计算
        :param expire_seconds: 缓存过期时间（秒）
        :param save_path: 相似度数据保存路径
        """
        # 检查缓存是否有效
        current_time = time.time()
        if (self.user_similarities is not None 
            and not force_recalculate 
            and current_time - self.last_similarity_update < expire_seconds):
            return self.user_similarities

        # 构建用户-模板行为矩阵
        all_users = list(self.user_behavior.keys())
        if len(all_users) < 2:
            self.user_similarities = {}
            return self.user_similarities

        user_template_matrix = []
        for user_id in all_users:
            behavior_vector = []
            for tid in self.template_ids:  # 使用同步后的模板ID列表
                behavior_vector.append(self.user_behavior[user_id].get(tid, 0))
            user_template_matrix.append(behavior_vector)

        # 计算余弦相似度
        user_similarity_matrix = cosine_similarity(user_template_matrix)

        # 构建相似度字典
        user_similarities = {}
        for i, target_user in enumerate(all_users):
            similar_users = {}
            for j, other_user in enumerate(all_users):
                if i != j:
                    similar_users[other_user] = float(user_similarity_matrix[i][j])
            user_similarities[target_user] = dict(
                sorted(similar_users.items(), key=lambda x: x[1], reverse=True)
            )

        # 更新缓存
        self.user_similarities = user_similarities
        self.last_similarity_update = current_time
        print(f"用户相似度计算完成，缓存已更新（共{len(all_users)}个用户）")

        # 保存到JSON文件
        try:
            with open(save_path, 'w', encoding='utf-8') as f:
                json.dump(user_similarities, f, ensure_ascii=False, indent=2)
            print(f"用户相似度数据已成功保存到: {save_path}")
        except Exception as e:
            print(f"保存用户相似度数据失败: {str(e)}")

        return self.user_similarities

    def after_behavior_updated(self):
        """用户行为更新后调用，标记缓存过期"""
        self.last_similarity_update = 0  # 下次调用时会重新计算


if __name__ == "__main__":
    # 初始化数据层
    print("\n" + "="*50)
    print("初始化数据准备层".center(40))
    print("="*50)
    data_layer = DataPreparationLayer()
    
    # 添加模板数据
    sample_templates = DataPreparationLayer.load_templates_from_json("/home/ai/Recommendation_system/database/template_data.json")
    
    print("\n" + "="*50)
    print("添加模板数据".center(40))
    print("="*50)
    data_layer.batch_add_templates(sample_templates)
    
    # 记录用户行为
    print("\n" + "="*50)
    print("记录用户行为".center(40))
    print("="*50)
    DataPreparationLayer.load_and_record_behavior(data_layer, "/home/ai/Recommendation_system/database/user_behavior_samples.json")


    # 保存数据
    print("\n" + "="*50)
    print("保存数据".center(40))
    print("="*50)
    data_layer.save_behavior_data()
    data_layer.export_template_vectors()
    
    # 打印状态报告
    print("\n" + "="*50)
    print("系统状态报告".center(40))
    print("="*50)
    print(f"• 模板数量：{data_layer.get_template_count()}")
    print(f"• 用户数量：{data_layer.get_user_count()}")
    print(f"• 模板1信息：{data_layer.get_template_info(1)}")
    
    # 保存所有用户相似度信息
    print("\n" + "="*50)
    print("保存所有用户相似度信息".center(40))
    print("="*50)
    user_similarities = data_layer._calculate_user_similarity(force_recalculate=True, save_path="/home/ai/Recommendation_system/user_similarities_data.json")
    print(f"• 用户相似度计算结果已保存到JSON文件，共{len(user_similarities)}个用户")
    


